from importlib import import_module
import pickle as pkl

import numpy as np

from utils.data_utils import *
from utils.training_utils import ModelCheckpoint
from config import Config
import argparse
import matplotlib.pyplot as plt


class ImgGenerator:
    def __init__(self, checkpt_path, config, char_map=None):
        """
        :param checkpt_path: Path of the model checkpoint file to be used
        :param config: Config with all the parameters to be used
        """

        self.config = config

        if char_map is None:
            with open(config.dict_file, 'r') as f:
                maps = {}
                lines = f.read().strip().split('\n')
                for i in range(len(lines)):
                    maps[lines[i]] = i + 1
                self.char_map = maps
        else:
            self.char_map = char_map

        # Model
        print(f'Model: {config.architecture}')
        model_type = import_module('models.' + self.config.architecture)
        create_model = getattr(model_type, 'create_model')
        self.model = create_model(self.config, self.char_map)
        # print(self.model, end='\n\n')
        self.model.to(self.config.device)
        self.model.eval()

        self.word_map = WordMap(self.char_map)

        # Load model weights
        self.model_checkpoint = ModelCheckpoint(config=self.config)
        self.model, _, _, _ = self.model_checkpoint.load(self.model, epoch=None, checkpoint_path=checkpt_path)

    def generate(self, random_num_imgs=5, word_list=None, z=None):
        """
        Returns images generated by the trained generator model

        :param random_num_imgs: Number of images to be randomly generated using lexicon (only valid if word_list=None)
        :param word_list: List of words for which images need to be generated
        :param z: Noise vector determining the style of the images to be generated (32 dimension vector)
        """

        if word_list is None:
            # Generate random images
            with torch.no_grad():
                self.model.forward_fake(z=None, b_size=random_num_imgs)
        else:
            encoded_words, _ = self.word_map.encode(word_list)
            with torch.no_grad():
                self.model.forward_fake(z=None, fake_y=encoded_words, b_size=len(word_list))

        word_labels_decoded = self.word_map.decode(self.model.fake_y.cpu().numpy())

        return self.model.fake_img.squeeze(1).cpu().numpy(), self.model.fake_y.cpu().numpy(), word_labels_decoded


if __name__ == "__main__":
    # Construct the argument parse and parse the arguments
    ap = argparse.ArgumentParser()
    ap.add_argument("-c", "--checkpt_path", required=True, type=str,
                    help="Path of the model checkpoint file to be used")
    # ap.add_argument("-m", "--char_map_path", required=True, type=str,
    #                 help="Path of the file with character mapping to be used")
    ap.add_argument("-n", "--num_imgs", required=False, type=int,
                    help="number of sample points")
    ap.add_argument("-w", "--word_list", required=False, nargs='+', default=[],
                    help="words for which images need to be generated")
    args = vars(ap.parse_args())
    checkpoint_path = args['checkpt_path']
    #char_map_path = args['char_map_path']
    num_imgs = args['num_imgs'] if args['num_imgs'] is not None else 5
    word_list = args['word_list'] if len(args['word_list']) > 0 else None

    # with open(f'{char_map_path}', 'rb') as f:
    #     char_map = pkl.load(f)
    char_map=None
    config = Config
    generator = ImgGenerator(checkpt_path=checkpoint_path, config=config, char_map=char_map)
    generated_imgs, _, word_labels = generator.generate(num_imgs, word_list)

    for label, img in zip(word_labels, generated_imgs):
        img=img*127.5+127.5
        img=img.astype(np.uint8)
        img=np.transpose(img,(1,2,0))
        plt.imshow(img,)
        plt.title(label)
        plt.show()
